import torch
import cv2
import numpy as np
from pathlib import Path
from models.common import DetectMultiBackend
from utils.general import (non_max_suppression, scale_boxes, check_img_size)
from utils.plots import Annotator, colors
from utils.torch_utils import select_device


class YOLOv5Detector:
    def __init__(self, weights='best-sw.pt', device='cpu', imgsz=640, conf_thres=0.25, iou_thres=0.45):
        """
        初始化 YOLOv5 目标检测器
        :param weights: 模型权重路径
        :param device: 运行设备 ('cpu' 或 'cuda:0')
        :param imgsz: 输入图像大小，可以是单个数值或 (height, width)
        :param conf_thres: 置信度阈值
        :param iou_thres: IoU 阈值
        """
        self.device = select_device(device)
        self.model = DetectMultiBackend(weights, device=self.device, dnn=False, fp16=False)

        # 确保 img_size 是一个 (H, W) 形式的元组
        self.imgsz = (imgsz, imgsz) if isinstance(imgsz, int) else tuple(imgsz)
        self.imgsz = check_img_size(self.imgsz, s=self.model.stride)

        self.conf_thres = conf_thres
        self.iou_thres = iou_thres
        self.names = self.model.names  # 类别名称
        self.model.warmup(imgsz=(1, 3, *self.imgsz))  # 预热模型

    def detect(self, img):
        """
        进行目标检测并返回带有标注的图像
        :param img: OpenCV 读取的 BGR 图像 (numpy.ndarray)
        :return: 带有检测框的图像 (numpy.ndarray)
        """
        assert isinstance(img, np.ndarray), "输入必须是 OpenCV 读取的 numpy.ndarray 格式"

        # 预处理图像
        img0 = img.copy()  # 保留原始图像
        img = cv2.resize(img, self.imgsz)  # 调整大小
        img = img[..., ::-1].transpose((2, 0, 1))  # BGR 转 RGB，通道 HWC -> CHW
        img = np.ascontiguousarray(img)  # 保持内存连续性
        img = torch.from_numpy(img).to(self.device).float() / 255.0  # 归一化
        img = img.unsqueeze(0) if len(img.shape) == 3 else img  # 添加 batch 维度

        # 推理
        with torch.no_grad():
            pred = self.model(img, augment=False, visualize=False)
            pred = non_max_suppression(pred, self.conf_thres, self.iou_thres, max_det=1000)

        # 处理检测结果
        annotator = Annotator(img0, line_width=3, example=str(self.names))
        for det in pred:
            if len(det):
                det[:, :4] = scale_boxes(img.shape[2:], det[:, :4], img0.shape).round()  # 调整检测框至原图尺寸
                for *xyxy, conf, cls in reversed(det):
                    # label = f"{self.names[int(cls)]} {conf:.2f}"
                    label = f"{self.names[int(cls)]}"
                    annotator.box_label(xyxy, label, color=colors(int(cls), True))

        return annotator.result()  # 返回带标注的图像

if __name__ == "__main__":
    detector = YOLOv5Detector(weights='best-sw.pt', device='cuda:0', imgsz=640)  # 初始化检测器
    img = cv2.imread('test.jpg')  # 读取测试图片
    result = detector.detect(img)  # 进行目标检测

    cv2.imshow("YOLOv5 Detection", result)  # 显示带有检测框的图片
    cv2.waitKey(0)
    cv2.destroyAllWindows()
